{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "分别合并数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<>:17: SyntaxWarning: invalid escape sequence '\\d'\n",
      "<>:17: SyntaxWarning: invalid escape sequence '\\d'\n",
      "C:\\Users\\dell\\AppData\\Local\\Temp\\ipykernel_81068\\2802676889.py:17: SyntaxWarning: invalid escape sequence '\\d'\n",
      "  result.to_csv('D:\\多任务选股_数据处理\\data_merged/data_vwap_merged.csv', index=False)      #需要修改\n"
     ]
    }
   ],
   "source": [
    "# 定义CSV文件所在的文件夹路径\n",
    "folder_path = r'D:\\多任务选股_数据处理\\raw_data\\vwap'  \n",
    "\n",
    "# 获取文件夹中所有的CSV文件\n",
    "csv_files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith('.csv')]\n",
    "\n",
    "dfs = []\n",
    "# 逐个读取CSV文件\n",
    "for file in csv_files:\n",
    "    df = pd.read_csv(file)\n",
    "    dfs.append(df)\n",
    "\n",
    "# 使用 pd.concat 自动对齐列（取列的并集），不同列自动填NaN\n",
    "result = pd.concat(dfs, axis=0, ignore_index=True)\n",
    "\n",
    "# 可选：保存到新CSV文件\n",
    "result.to_csv('D:\\多任务选股_数据处理\\data_merged/data_vwap_merged.csv', index=False)      "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "日期转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_date(date_str):\n",
    "    possible_formats = ['%Y%m%d', '%Y/%m/%d']\n",
    "    for fmt in possible_formats:\n",
    "        try:\n",
    "            date_obj = datetime.strptime(date_str, fmt)\n",
    "            return date_obj.strftime('%Y-%m-%d')\n",
    "        except ValueError:\n",
    "            continue\n",
    "    # 如果输入已经是 YYYYMMDD 格式或者无法解析，直接返回原字符串\n",
    "    return date_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "industry_df=pd.read_csv('industry_data.csv')\n",
    "industry_df['entry_dt'] = industry_df['entry_dt'].astype(str)\n",
    "industry_df['entry_dt'] = industry_df['entry_dt'].apply(convert_date)\n",
    "industry_df['remove_dt'] = industry_df['remove_dt'].astype(str)\n",
    "# 去除小数点及后面的零\n",
    "industry_df['remove_dt'] = industry_df['remove_dt'].str.rstrip('.0')\n",
    "industry_df['remove_dt'] = industry_df['remove_dt'].apply(convert_date)\n",
    "industry_df.to_csv('D:/多任务选股_数据处理/industry_df.csv', index=False)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "剔除ST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "总共找到 645495 个 ST 股票标记\n"
     ]
    }
   ],
   "source": [
    "st = pd.read_csv('isst.csv',index_col=0, parse_dates=True) \n",
    "#遍历 isst 表，找到所有值为1的（即 ST 股票）\n",
    "# 用 stack 将 DataFrame 转为长格式，方便筛选\n",
    "st_stacked = st.stack()  # MultiIndex: (日期, 股票代码), 值\n",
    "st_positions = st_stacked[st_stacked == 1]  # 取值为 1 的位置\n",
    "print(f\"总共找到 {len(st_positions)} 个 ST 股票标记\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dt                   \n",
      "2005-12-30  000005.SZ    1.0\n",
      "            000017.SZ    1.0\n",
      "            000020.SZ    1.0\n",
      "            000025.SZ    1.0\n",
      "            000030.SZ    1.0\n",
      "                        ... \n",
      "2024-12-17  603879.SH    1.0\n",
      "            603959.SH    1.0\n",
      "            603963.SH    1.0\n",
      "            688282.SH    1.0\n",
      "            688287.SH    1.0\n",
      "Length: 645495, dtype: float64\n"
     ]
    }
   ],
   "source": [
    "print(st_positions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "已完成所有替换\n"
     ]
    }
   ],
   "source": [
    "list=['amt','close', 'high', 'low', 'open', 'turn', 'volume', 'vwap']\n",
    "for i in list:\n",
    "    data_table = pd.read_csv(f'data_merged/data_{i}_merged.csv',index_col=0, parse_dates=True)\n",
    "    #替换 data_table 中相应的值为 NaN\n",
    "    for date, stock_code in st_positions.index:\n",
    "        if date in data_table.index and stock_code in data_table.columns:\n",
    "            data_table.loc[date, stock_code] = np.nan  \n",
    "    data_table.to_csv(f'data_remove/data_{i}_remove_st.csv')\n",
    "print(\"已完成所有替换\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "剔除上市不足三个月的股票"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                 证券代码  证券简称\n",
      "上市日期                       \n",
      "1991-04-03  000001.SZ  平安银行\n",
      "1991-01-29  000002.SZ   万科A\n",
      "1990-12-01  000004.SZ  国华网安\n",
      "1992-04-27  000006.SZ  深振业A\n",
      "1992-04-13  000007.SZ   全新好\n"
     ]
    }
   ],
   "source": [
    "start_date_data = pd.read_excel('start_date_data.xlsx')\n",
    "start_date_data = start_date_data.set_index('上市日期')\n",
    "print(start_date_data.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<>:4: SyntaxWarning: invalid escape sequence '\\d'\n",
      "<>:13: SyntaxWarning: invalid escape sequence '\\d'\n",
      "<>:4: SyntaxWarning: invalid escape sequence '\\d'\n",
      "<>:13: SyntaxWarning: invalid escape sequence '\\d'\n",
      "C:\\Users\\dell\\AppData\\Local\\Temp\\ipykernel_75312\\1583391707.py:4: SyntaxWarning: invalid escape sequence '\\d'\n",
      "  stock_data = pd.read_csv(f'D:\\多任务选股_数据处理\\data_remove\\data_{i}_remove_st.csv', index_col=0)\n",
      "C:\\Users\\dell\\AppData\\Local\\Temp\\ipykernel_75312\\1583391707.py:13: SyntaxWarning: invalid escape sequence '\\d'\n",
      "  stock_data.to_csv(f'D:\\多任务选股_数据处理\\data_processed\\data_{i}_processed.csv')\n"
     ]
    }
   ],
   "source": [
    "start_date_data.index = pd.to_datetime(start_date_data.index)\n",
    "for i in list:\n",
    "    stock_data = pd.read_csv(f'D:\\多任务选股_数据处理\\data_remove\\data_{i}_remove_st.csv', index_col=0)\n",
    "    stock_data.index = pd.to_datetime(stock_data.index)\n",
    "    for stock_code in stock_data.columns:\n",
    "        # 获取该股票的上市时间\n",
    "        listing_date = start_date_data.loc[start_date_data['证券代码'] == stock_code].index[0]\n",
    "        # 计算上市时间后 3 个月的日期\n",
    "        three_months_after = listing_date + pd.DateOffset(months=3)\n",
    "        # 将上市时间未满 3 个月的股票数据替换为 NaN\n",
    "        stock_data.loc[(stock_data.index < three_months_after), stock_code] = float('nan')\n",
    "    stock_data.to_csv(f'D:\\多任务选股_数据处理\\data_processed\\data_{i}_processed.csv')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
